import re
import regex
from collections import Counter

from typing import Optional

ERROR_BASE = re.compile("[^ATGCNatgcn]")
GEN_CODE = {
    'ATA':'I', 'ATC':'I', 'ATT':'I', 'ATG':'M',
    'ACA':'T', 'ACC':'T', 'ACG':'T', 'ACT':'T',
    'AAC':'N', 'AAT':'N', 'AAA':'K', 'AAG':'K',
    'AGC':'S', 'AGT':'S', 'AGA':'R', 'AGG':'R',
    'CTA':'L', 'CTC':'L', 'CTG':'L', 'CTT':'L',
    'CCA':'P', 'CCC':'P', 'CCG':'P', 'CCT':'P',
    'CAC':'H', 'CAT':'H', 'CAA':'Q', 'CAG':'Q',
    'CGA':'R', 'CGC':'R', 'CGG':'R', 'CGT':'R',
    'GTA':'V', 'GTC':'V', 'GTG':'V', 'GTT':'V',
    'GCA':'A', 'GCC':'A', 'GCG':'A', 'GCT':'A',
    'GAC':'D', 'GAT':'D', 'GAA':'E', 'GAG':'E',
    'GGA':'G', 'GGC':'G', 'GGG':'G', 'GGT':'G',
    'TCA':'S', 'TCC':'S', 'TCG':'S', 'TCT':'S',
    'TTC':'F', 'TTT':'F', 'TTA':'L', 'TTG':'L',
    'TAC':'Y', 'TAT':'Y', 'TAA':'_', 'TAG':'_',
    'TGC':'C', 'TGT':'C', 'TGA':'_', 'TGG':'W'}

def validate_dna(seq:str) -> None:
    match_iter = re.finditer(ERROR_BASE,seq)
    for error in match_iter:
        raise ValueError(
            "{} at {} index is unsupported base".format(
                error.group()[0],error.span()[0]))

def frequency(seq:str) -> Counter:
    return Counter(seq)

def gc_content(seq:str) -> float:
    counter = frequency(seq)
    num_gc = 0
    for base in "GCag":
        num_gc += counter[base]
    return num_gc / len(seq)

def gc_content_subseq(seq:str,k:Optional[int]=100) -> list[float]:
    gc_sub = list()
    for i in range(0,len(seq)-k+1,k):
        gc_sub.append(
            gc_content(seq[i:i+k])
        )
    return gc_sub

def transcription(seq:str) -> str:
    validate_dna(seq)
    return seq.upper().replace("T","U")

def reverse_complement(seq:str) -> str:
    '''
    https://blog.csdn.net/lotusng/article/details/103315213
    '''
    validate_dna(seq)
    trantab = str.maketrans('ACGTacgtRYMKrymkVBHDvbhd', 'TGCAtgcaYRKMyrkmBVDHbvdh')
    seq = seq.translate(trantab)
    return seq[::-1]

def translate_codon(codon) -> str:
    return GEN_CODE.get(codon,'X')

def translate_seq(seq:str,ini_pos=0) -> str:
    validate_dna(seq)
    seq = seq.upper()
    return ''.join((translate_codon(seq[pos:pos+3]) \
                    for pos in range(ini_pos,len(seq)-2,3)))

def codon_usage(seq:str) -> Counter[str:float]:
    aa_counter = Counter(translate_seq(seq))
    num_aa = sum(aa_counter.values())
    for k,v in aa_counter.items():
        aa_counter[k] = v/num_aa
    return aa_counter

def reading_frames(seq:str) -> list[str]:
    '''
    computes the translation of the six different reading frames
    return in list shape: [+0,+1,+2,-0,-1,-2]
    '''
    validate_dna(seq)
    res = list()
    for _ in range(2):
        for ini_pos in (0,1,2):
            res.append(translate_seq(seq,ini_pos))
        seq = reverse_complement(seq)
    return res

def all_proteins_rf(aa_seq) -> list[str]:
    return regex.findall(r"(M\w*?)_",aa_seq,overlapped=True)

def all_orfs(seq:str,**karg) -> list[str]:
    '''
    all putative proteins in all reading frames
    '''
    min_size = karg.get("min_size",0)

    reading_frame_seq = reading_frames(seq)
    proteins = list()
    for seq in reading_frame_seq:
        for protein in all_proteins_rf(seq):
            if len(protein) > min_size:
                proteins.append(protein)
    
    if isinstance(karg.get("min_size",True),bool):
        return proteins
    else:
        return sorted(proteins,key=lambda s:len(s),reverse=True)

def all_orfs_ord(seq:str,min_size:int=0):
    return all_orfs(seq,min_size=min_size)


if __name__ == "__main__":
    # upper_dna_seq("ATCgaxaGA")
    print(gc_content("ATCTCTGATCGCGCTAGA"))
    print(gc_content_subseq("ATCTCTGATCGCGCTAGA",5))
    print(transcription("ATCTCTGATCGCGCTAGA"))
    print(reverse_complement("ATCTCTGATCGCGCTAGA"))
    print(translate_seq("ATCTCTGATCGCGCTAGA"))
    print(codon_usage("ATCTCTGATCGCGCTAGA"))
    print(reading_frames("ATCTCTGATCGCGCTAGA"))
    print(all_proteins_rf("ASDSFMASDFAS_SAFMDSFMASD_S"))
    print(all_orfs("TCGATCAGCATGCTTAGCTAGCAGCTGACATGCt"))
    print(all_orfs_ord("TCGATCAGCATGCTTAGCTAGCAGCTGACATGCt"))
    pass
